from sqlalchemy import Column, String, Integer
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

from init.mysql import engine

Base = declarative_base()


class ConceptInfo(Base):
    __tablename__ = 'concept_info'

    code = Column(String(20), primary_key=True)
    name = Column(String(100))
    type1 = Column(Integer)
    type2 = Column(Integer)
    level = Column(Integer)
    pcode = Column(String(20))
    pname = Column(String(100))
    isleaf = Column(Integer)

    def __init__(self, code=None, name=None, type1=None, type2=None, level=None, pcode=None, pname=None, isleaf=None):
        self.code = code
        self.name = name
        self.type1 = type1
        self.type2 = type2
        self.level = level
        self.pcode = pcode
        self.pname = pname
        self.isleaf = isleaf


Base.metadata.create_all(engine)


def batch_insert_or_update(data_list):
    session = sessionmaker(bind=engine)()
    try:
        batch_size = 1000
        i = 0
        for data in data_list:
            # 查询是否已存在相同的记录
            existing_record = session.query(ConceptInfo).filter_by(code=data['code']).first()
            if existing_record:
                # 如果存在，更新记录
                existing_record.code = data['code']
                existing_record.name = data['name']
                existing_record.type1 = data['type1']
                existing_record.type2 = data['type2']
                existing_record.level = data['level']
                existing_record.pcode = data['pcode']
                existing_record.pname = data['pname']
                existing_record.isleaf = data['isleaf']
            else:
                # 如果不存在，插入新记录
                new_record = ConceptInfo(code=data['code'], name=data['name'], type1=data['type1'], type2=data['type2'],
                                         level=data['level'], pcode=data['pcode'], pname=data['pname'],
                                         isleaf=data['isleaf'])
                session.add(new_record)
            i += 1
            if i % batch_size == 0:
                i = 0
                # 提交事务
                session.commit()
    except Exception as e:
        print(f"Error: {e}")
        session.rollback()
    finally:
        session.close()
